package com.ocean.web.util.nccalculate;
import ucar.ma2.*;
import ucar.nc2.*;

import java.io.IOException;
import java.util.*;

public class CalculateWindSpeedAndDirXBDQSZ {

    public static void main(String[] args) {
        // 输入文件路径（之前生成的包含U10、V10的NC文件）
        String inputNcPath = "E:\\NMF_BEN_SA_NWPACDT_2025102800_168h_ATM_cropped.nc";
        // 输出文件路径（添加wspd10和wdir10后的新文件）
        String outputNcPath = "E:\\NMF_BEN_SA_NWPACDT_2025102800_168h_ATM_calculate.nc";
        CalculateWindSpeedAndDir(inputNcPath,outputNcPath);
    }

    /**
     * 计算方法
     * @param inputNcPath 计算的文件
     * @param outputNcPath 计算完的文件
     */
    public static boolean CalculateWindSpeedAndDir(String inputNcPath, String outputNcPath){
        try {
            NetcdfFile inputNc = NetcdfFile.open(inputNcPath);
            System.out.println("成功打开输入文件：" + inputNcPath);

            // 读取U10和V10变量
            Variable u10Var = inputNc.findVariable("U10");
            Variable v10Var = inputNc.findVariable("V10");
            if (u10Var == null || v10Var == null) {
                throw new RuntimeException("输入文件中未找到U10或V10变量");
            }

            // 校验维度一致性
            if (!u10Var.getDimensions().equals(v10Var.getDimensions())) {
                throw new RuntimeException("U10和V10变量维度不一致");
            }

            // 读取数据并计算风速、风向
            Array u10Array = u10Var.read();
            Array v10Array = v10Var.read();
            float[] u10Data = (float[]) u10Array.copyTo1DJavaArray();
            float[] v10Data = (float[]) v10Array.copyTo1DJavaArray();

            int dataSize = u10Data.length;
            float[] wspd10Data = new float[dataSize];
            float[] wdir10Data = new float[dataSize];
            for (int i = 0; i < dataSize; i++) {
                float u = u10Data[i];
                float v = v10Data[i];
                wspd10Data[i] = (float) Math.sqrt(u * u + v * v); // 风速
                double rad = Math.atan2(u, v);
                wdir10Data[i] = (float) (rad * 180 / Math.PI + 180); // 风向
            }

            // 修复：将outputPath改为outputNcPath，与定义的变量名一致
            createOutputNcFile(inputNc, u10Var, wspd10Data, wdir10Data, outputNcPath);

            inputNc.close();
            System.out.println("计算完成！新文件：" + outputNcPath);
            return true;
        } catch (Exception e) {
            System.err.println("处理失败：" + e.getMessage());
            e.printStackTrace();
            return false;
        }
    }


    /**
     * 创建输出文件（兼容低版本NetCDF-Java）
     */
    private static void createOutputNcFile(NetcdfFile inputNc, Variable u10Var,
                                           float[] wspd10Data, float[] wdir10Data, String outputPath) throws IOException, InvalidRangeException {

        NetcdfFileWriter writer = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf4, outputPath);

        try {
            Group rootGroup = null; // 低版本用null表示根组

            // 1. 收集并去重维度
            Set<String> dimNameSet = new HashSet<>();
            List<Dimension> allRequiredDims = new ArrayList<>();
            List<Variable> inputVars = inputNc.getVariables();
            for (Variable var : inputVars) {
                for (Dimension dim : var.getDimensions()) {
                    String dimName = dim.getShortName();
                    if (!dimNameSet.contains(dimName)) {
                        dimNameSet.add(dimName);
                        allRequiredDims.add(dim);
                    }
                }
            }

            // 2. 创建维度
            Map<String, Dimension> createdDims = new HashMap<>();
            for (Dimension dim : allRequiredDims) {
                String dimName = dim.getShortName();
                Dimension newDim = writer.addDimension(rootGroup, dimName, dim.getLength());
                createdDims.put(dimName, newDim);
            }

            // 3. 创建原变量
            Map<String, Variable> createdVars = new HashMap<>();
            for (Variable var : inputVars) {
                String varName = var.getShortName();
                if (createdVars.containsKey(varName)) continue;

                List<Dimension> varDims = new ArrayList<>();
                for (Dimension dim : var.getDimensions()) {
                    varDims.add(createdDims.get(dim.getShortName()));
                }

                Variable newVar = writer.addVariable(rootGroup, varName, var.getDataType(), varDims);
                copyVariableAttributes(var, newVar);
                createdVars.put(varName, newVar);
            }

            // 4. 添加新变量
            List<Dimension> dataDims = new ArrayList<>();
            for (Dimension dim : u10Var.getDimensions()) {
                dataDims.add(createdDims.get(dim.getShortName()));
            }

            Variable wspd10Var = writer.addVariable(rootGroup, "wspd10", DataType.FLOAT, dataDims);
            wspd10Var.addAttribute(new Attribute("long_name", "10米高度风速"));
            wspd10Var.addAttribute(new Attribute("units", "m/s"));

            Variable wdir10Var = writer.addVariable(rootGroup, "wdir10", DataType.FLOAT, dataDims);
            wdir10Var.addAttribute(new Attribute("long_name", "10米高度风向"));
            wdir10Var.addAttribute(new Attribute("units", "degree"));

            // 5. 复制全局属性
            for (Attribute attr : inputNc.getGlobalAttributes()) {
                writer.addGlobalAttribute(attr);
            }

            // 6. 写入数据
            writer.create();
            for (Variable var : inputVars) {
                Variable destVar = createdVars.get(var.getShortName());
                if (destVar != null) {
                    writer.write(destVar, var.read());
                }
            }

            int[] shape = u10Var.getShape();
            writer.write(wspd10Var, Array.factory(DataType.FLOAT, shape, wspd10Data));
            writer.write(wdir10Var, Array.factory(DataType.FLOAT, shape, wdir10Data));

            System.out.println("新变量添加成功：wspd10、wdir10");

        } finally {
            writer.close();
        }
    }

    private static void copyVariableAttributes(Variable srcVar, Variable destVar) {
        for (Attribute attr : srcVar.getAttributes()) {
            destVar.addAttribute(attr);
        }
    }
}